{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EgiF12Hf1Dhs"
   },
   "source": [
    "This notebook provides examples to go along with the [textbook](http://manipulation.csail.mit.edu/trajectories.html).  I recommend having both windows open, side-by-side!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "eeMrMI0-1Dhu"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from pydrake.all import (\n",
    "    AddMultibodyPlantSceneGraph,\n",
    "    DiagramBuilder,\n",
    "    InverseKinematics,\n",
    "    MeshcatVisualizer,\n",
    "    MeshcatVisualizerParams,\n",
    "    Parser,\n",
    "    RigidTransform,\n",
    "    Solve,\n",
    "    Sphere,\n",
    "    StartMeshcat,\n",
    ")\n",
    "\n",
    "from manipulation import running_as_notebook\n",
    "from manipulation.meshcat_utils import plot_mathematical_program\n",
    "from manipulation.scenarios import AddShape, AddTwoLinkIiwa, AddWsg\n",
    "from manipulation.utils import ConfigureParser"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start the visualizer.\n",
    "meshcat = StartMeshcat()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Visualizing the costs and constraints\n",
    "\n",
    "Here is another view.  Notice that at the optimal solution of the iiwa reaching into the shelf, the last joint was almost zero.  I've gone ahead and welded it to zero, so that we are now down to just a two degree of freedom IK problem.  Now we can plot the entire cost and constraint landscape.  Let's do it.\n",
    "\n",
    "There is a lot going on in the second meshcat window.  Use the controls to turn on and off different costs/constraints.  The constraints are blue where they are feasible and red where they are infeasible.  Which constraint is the horribly ugly one?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def draw_ik_prog(zoom=True):\n",
    "    builder = DiagramBuilder()\n",
    "\n",
    "    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.001)\n",
    "    iiwa = AddTwoLinkIiwa(plant)\n",
    "    wsg = AddWsg(plant, iiwa, roll=0.0, welded=True)\n",
    "    sphere = AddShape(plant, Sphere(0.02), \"sphere\")\n",
    "    X_WO = RigidTransform([0.6, 0, 0.65])\n",
    "    plant.WeldFrames(plant.world_frame(), plant.GetFrameByName(\"sphere\"), X_WO)\n",
    "\n",
    "    parser = Parser(plant)\n",
    "    ConfigureParser(parser)\n",
    "    bin = parser.AddModelsFromUrl(\"package://manipulation/shelves.sdf\")[0]\n",
    "    plant.WeldFrames(\n",
    "        plant.world_frame(),\n",
    "        plant.GetFrameByName(\"shelves_body\", bin),\n",
    "        RigidTransform([0.6, 0, 0.4]),\n",
    "    )\n",
    "\n",
    "    plant.Finalize()\n",
    "\n",
    "    #    visualizer = MeshcatVisualizer.AddToBuilder(\n",
    "    #        builder,\n",
    "    #        scene_graph,\n",
    "    #        meshcat,\n",
    "    #        MeshcatVisualizerParams(delete_prefix_initialization_event=False))\n",
    "\n",
    "    diagram = builder.Build()\n",
    "    context = diagram.CreateDefaultContext()\n",
    "    plant_context = plant.GetMyContextFromRoot(context)\n",
    "\n",
    "    q0 = plant.GetPositions(plant_context)\n",
    "    gripper_frame = plant.GetFrameByName(\"body\", wsg)\n",
    "    sphere_frame = plant.GetFrameByName(\"sphere\", sphere)\n",
    "\n",
    "    ik = InverseKinematics(plant, plant_context)\n",
    "    collision_constraint = ik.AddMinimumDistanceLowerBoundConstraint(\n",
    "        0.001, 0.1\n",
    "    )\n",
    "    grasp_constraint = ik.AddPositionConstraint(\n",
    "        gripper_frame,\n",
    "        [0, 0.1, 0],\n",
    "        sphere_frame,\n",
    "        [-0.001, -0.001, -0.001],\n",
    "        [0.001, 0.001, 0.001],\n",
    "    )\n",
    "\n",
    "    prog = ik.get_mutable_prog()\n",
    "    q = ik.q()\n",
    "    prog.AddQuadraticErrorCost(np.identity(len(q)), q0, q)\n",
    "    prog.SetInitialGuess(q, q0)\n",
    "    result = Solve(ik.prog())\n",
    "    if not result.is_success():\n",
    "        print(\"IK failed\")\n",
    "\n",
    "    diagram.ForcedPublish(context)\n",
    "\n",
    "    meshcat.Delete()\n",
    "    meshcat.SetProperty(\"/Background\", \"visible\", False)\n",
    "    if zoom:\n",
    "        qstar = result.GetSolution(q)\n",
    "        X, Y = np.meshgrid(\n",
    "            np.linspace(qstar[0] - 0.2, qstar[0] + 0.2, 75),\n",
    "            np.linspace(qstar[1] - 0.2, qstar[1] + 0.2, 75),\n",
    "        )\n",
    "        point_size = 0.01\n",
    "    else:\n",
    "        low = plant.GetPositionLowerLimits()\n",
    "        up = plant.GetPositionUpperLimits()\n",
    "        X, Y = np.meshgrid(\n",
    "            np.linspace(low[0], up[0], 175), np.linspace(low[1], up[1], 175)\n",
    "        )\n",
    "        point_size = 0.05\n",
    "    plot_mathematical_program(\n",
    "        meshcat, \"ik\", prog, X, Y, result, point_size=point_size\n",
    "    )\n",
    "\n",
    "\n",
    "draw_ik_prog(zoom=True)"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "Robotic Manipulation - Motion Planning.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3.8.10 64-bit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.6"
  },
  "vscode": {
   "interpreter": {
    "hash": "e7370f93d1d0cde622a1f8e1c04877d8463912d04d973331ad4851f04de6915a"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
